#!/usr/bin/python
# -*- coding:utf-8 -*-
# @author  : 龍皓
# @time    : 2024/12/21 11:23
# @function: 视频工具类
# @version : 1.0.0
import re
import time
import cv2
import os
import shutil
import easyocr
from glob import glob
from pathlib import Path
import numpy as np
from PIL import Image
from concurrent.futures import ThreadPoolExecutor



class DouyinUtil:

    def extract_last_frame(self, video_path, output_dir):
        # 使用绝对路径来避免路径问题
        video_path = os.path.abspath(video_path)

        # 打开视频文件
        cap = cv2.VideoCapture(video_path)

        # 检查视频是否成功打开
        if not cap.isOpened():
            print(f"无法打开视频文件: {video_path}")
            return None

        # 获取视频的总帧数
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        # print(f"视频总帧数: {total_frames}")

        if total_frames == 0:
            print("视频没有帧内容")
            cap.release()
            return None

        # 逐帧读取，直到最后一帧
        for _ in range(total_frames - 2):
            ret, frame = cap.read()
            if not ret:
                break

        # 读取最后一帧
        ret, frame = cap.read()

        if ret:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            # 保存图片
            base_filename = Path(video_path).stem
            filename = os.path.normpath(
                os.path.join(output_dir, f"last_frame_{self.clean_folder_name(base_filename)}.jpg"))

            # 打印调试信息，检查文件路径
            # print(f"保存最后一帧：{filename}")

            # 使用 Pillow 保存文件
            pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))  # 转为RGB模式
            pil_image.save(str(filename), format='JPEG')  # 使用Pillow保存

            # print(f"图片已保存到: {filename}")
            if not os.path.exists(filename):  # 确认文件是否保存成功
                print(f"保存图片失败：{filename}")
            cap.release()
            return filename
        else:
            print(f"无法读取最后一帧：{video_path}")
            cap.release()
            return None

    def ocr_recognize_text(self, image_path):

        if not image_path or not os.path.exists(image_path):
            print("无效的图像路径或文件不存在:", image_path)
            return None

        # 使用 PIL 打开图像
        try:
            image = Image.open(image_path)
        except Exception as e:
            print(f"无法读取图像文件: {image_path}, 错误: {e}")
            return None

            # 检查图像是否成功加载
        if image is None:
            print("图像为空，无法继续")
            return None

            # 转为灰度图像并转换为 NumPy 数组
        gray_image = image.convert("L")
        gray_image_np = np.array(gray_image)

        # OCR识别
        reader = easyocr.Reader(['en', 'ch_sim'], gpu=True)
        result = reader.readtext(gray_image_np)
        text = ""

        for detection in result:
            text += detection[1] + " "
        text = text.strip()

        # 提取 "抖音号: xxx" 这部分，并去掉最后的 "抖音"
        match = re.search(r'抖音号[:：]\s*([\w_]+)', text)
        if match:
            # 获取抖音号后面的内容
            extracted_text = match.group(1)  # 获取用户名
            print(extracted_text)
            return extracted_text
        return text.strip()  # 如果没有匹配，返回原始识别文本

    def create_folder_and_save(self, folder_base, video_path, text, last_frame_path):
        # 生成基于文字的文件夹名称
        folder_name = self.clean_folder_name(text) if text else "no_text"
        folder_path = os.path.join(folder_base, folder_name)

        # 只在文件夹不存在时创建文件夹
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)

        # 保存视频到新文件夹
        video_dest = os.path.join(folder_path, Path(video_path).name)
        self.safe_move(video_path, video_dest)  # 移动视频文件到新文件夹

        # 如果存在最后一帧图片，移动它到新文件夹
        if last_frame_path:
            # 构建目标文件路径
            last_frame_dest = os.path.join(folder_path, Path(last_frame_path).name)

            # 如果目标文件已经存在，修改文件名以避免覆盖
            if os.path.exists(last_frame_dest):
                base, ext = os.path.splitext(last_frame_dest)
                timestamp = int(time.time())  # 使用时间戳或你也可以使用递增数字
                last_frame_dest = f"{base}_{timestamp}{ext}"
            try:
                # 显式指定目标文件名的编码方式，避免乱码
                last_frame_dest = last_frame_dest.encode('utf-8', 'ignore').decode('utf-8')
                self.safe_move(last_frame_path, last_frame_dest)  # 移动文件
            except Exception as e:
                print(f"移动最后一帧图片时出现错误: {e}")

        print(f"视频已保存到: {video_dest}")
        return folder_path

    def get_video_files(self, folder_path):
        # 获取文件夹中所有视频文件
        video_files = glob(os.path.join(folder_path, '*.[mM][pP]4')) + \
                      glob(os.path.join(folder_path, '*.[aA][vV][iI]')) + \
                      glob(os.path.join(folder_path, '*.[mM][oO][vV]'))
        return video_files

    def clean_folder_name(self, text):
        # print(f"原始文件名: {text}")  # 打印原始文件名

        # 移除 Windows 不允许的字符（包括: \ / : * ? " < > |）
        text = re.sub(r'[\\/*?:"<>|]', "", text)  # 清除非法字符
        # print(f"清理非法字符后的文件名: {text}")

        # 替换空格为下划线，避免路径中出现空格
        text = re.sub(r'\s+', '_', text)
        # print(f"替换空格后的文件名: {text}")

        # 对特殊符号进行替换（如#，防止乱码）
        text = re.sub(r'[^\w\s-]', '_', text)
        # print(f"去除特殊字符后的文件名: {text}")

        # 限制文件夹名长度 (Windows 最长文件路径为 260，文件夹名一般不超过 100 字符)
        text = text[:100]

        return text

    def safe_move(self, src, dst):
        """
            安全地将文件从 src 移动到 dst。支持同分区和跨分区。
            """
        try:
            os.rename(src, dst)  # 同分区时，os.rename 高效处理
        except OSError as e:
            if "系统无法将文件移到不同的磁盘驱动器" in str(e):
                # 如果跨分区，改为 shutil.copy2 + os.remove
                shutil.copy2(src, dst)  # 复制文件
                os.remove(src)  # 删除原始文件
            else:
                raise e  # 抛出其他异常

def process_video(vu, video_file, output_dir):
    # 提取每个视频的最后一帧
    last_frame = vu.extract_last_frame(video_file, output_dir)

    if last_frame:
        # 使用OCR识别最后一帧中的文字
        recognized_text = vu.ocr_recognize_text(last_frame)
        print(f"视频: {video_file}, 识别出的文字: {recognized_text}")

        # 根据识别的文字创建文件夹，并将视频移动到该文件夹
        vu.create_folder_and_save(output_dir, video_file, recognized_text, last_frame)

if __name__ == '__main__':
    # 创建VideoUtil对象
    du = DouyinUtil()

    folder_path = "../media"  # 视频文件夹路径
    output_dir = "../outfile"  # 存储最后一帧的目录

    # video_path = '../media/share_03ebc6456d6d2ffe02a0a1ac508ab070.mp4'
    # cap = cv2.VideoCapture(video_path)
    # if not cap.isOpened():
    #     print("无法打开视频文件")
    # else:
    #     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    #     print(f"总帧数: {total_frames}")
    #
    #     for i in range(total_frames - 2):  # 逐帧读取直到倒数第二帧
    #         ret, frame = cap.read()
    #         if not ret:
    #             print(f"无法读取第{i + 1}帧")
    #             break
    #         print(f"读取第{i + 1}帧")
    #
    #     ret, frame = cap.read()
    #     print(ret)
    #     if ret:
    #         # 保存图片
    #         base_filename = Path(video_path).stem
    #         filename = os.path.normpath(
    #             os.path.join(output_dir, f"last_frame.jpg"))
    #         print(filename)
    #         # 使用 Pillow 保存文件
    #         pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))  # 转为RGB模式
    #         pil_image.save(str(filename), format='JPEG')  # 使用Pillow保存
    #
    #         # print(f"图片已保存到: {filename}")
    #         if not os.path.exists(filename):  # 确认文件是否保存成功
    #             print(f"保存图片失败：{filename}")
    #         cap.release()

    # 获取文件夹中的所有视频文件
    video_files = du.get_video_files(folder_path)
    # print("正在提取: ", video_files)

    # 使用ThreadPoolExecutor进行多线程处理
    with ThreadPoolExecutor(max_workers=4) as executor:
        for video_file in video_files:
            executor.submit(process_video, du, video_file, output_dir)
